#include "common/memory.h"

#include <cstddef>
#include <cxxabi.h>
#include <execinfo.h>
#include <new>
#include <string>

extern "C" {
#include <postgres.h>

#include <utils/memutils.h>
}

void* operator new(std::size_t size) {
  if (auto* ptr = palloc_extended(size, MCXT_ALLOC_NO_OOM); ptr != nullptr) {
    optimizer::AddCheck(ptr, size);
    return ptr;
  }
  throw std::bad_alloc();
}

void* operator new(std::size_t size, std::align_val_t align) {
  if (auto* ptr = palloc_aligned(size, static_cast<size_t>(align), MCXT_ALLOC_NO_OOM); ptr != nullptr) {
    optimizer::AddCheck(ptr, size);
    return ptr;
  }
  throw std::bad_alloc();
}

void* operator new[](std::size_t size) {
  if (auto* ptr = palloc_extended(size, MCXT_ALLOC_NO_OOM); ptr != nullptr) {
    optimizer::AddCheck(ptr, size);
    return ptr;
  }
  throw std::bad_alloc();
}

void* operator new[](std::size_t size, std::align_val_t align) {
  if (auto* ptr = palloc_aligned(size, static_cast<size_t>(align), MCXT_ALLOC_NO_OOM); ptr != nullptr) {
    optimizer::AddCheck(ptr, size);
    return ptr;
  }
  throw std::bad_alloc();
}

void* operator new(std::size_t size, const std::nothrow_t& /*unused*/) noexcept {
  auto* ptr = palloc_extended(size, MCXT_ALLOC_NO_OOM);
  optimizer::AddCheck(ptr, size);
  return ptr;
}

void* operator new[](std::size_t size, const std::nothrow_t& /*unused*/) noexcept {
  auto* ptr = palloc_extended(size, MCXT_ALLOC_NO_OOM);
  optimizer::AddCheck(ptr, size);
  return ptr;
}

void* operator new(std::size_t size, std::align_val_t align, const std::nothrow_t& /*unused*/) noexcept {
  auto* ptr = palloc_aligned(size, static_cast<size_t>(align), MCXT_ALLOC_NO_OOM);
  optimizer::AddCheck(ptr, size);
  return ptr;
}

void* operator new[](std::size_t size, std::align_val_t align, const std::nothrow_t& /*unused*/) noexcept {
  auto* ptr = palloc_aligned(size, static_cast<size_t>(align), MCXT_ALLOC_NO_OOM);
  optimizer::AddCheck(ptr, size);
  return ptr;
}

void operator delete(void* ptr) noexcept {
  optimizer::DeleCheck(ptr);
  pfree(ptr);
}

void operator delete(void* ptr, std::align_val_t align) noexcept {
  optimizer::DeleCheck(ptr);
  pfree(ptr);
}

void operator delete[](void* ptr) noexcept {
  optimizer::DeleCheck(ptr);
  pfree(ptr);
}

void operator delete[](void* ptr, std::align_val_t align) noexcept {
  optimizer::DeleCheck(ptr);
  pfree(ptr);
}

void operator delete(void* ptr, std::size_t size) noexcept {
  optimizer::DeleCheck(ptr);
  pfree(ptr);
}

void operator delete(void* ptr, std::size_t size, std::align_val_t align) noexcept {
  optimizer::DeleCheck(ptr);
  pfree(ptr);
}

void operator delete[](void* ptr, std::size_t size) noexcept {
  optimizer::DeleCheck(ptr);
  pfree(ptr);
}

void operator delete[](void* ptr, std::size_t size, std::align_val_t align) noexcept {
  optimizer::DeleCheck(ptr);
  pfree(ptr);
}

namespace optimizer {

#if PLAN_DEBUG

UNUSED_ATTRIBUTE static std::string Addr2sym(void* addr) {
  if (addr == nullptr) {
    return {"null"};
  }
  char** strings = backtrace_symbols(&addr, 1);
  if (strings == nullptr) {
    return {"???"};
  }
  std::string ret{strings[0]};
  free(static_cast<void*>(strings));  // NOLINT
  auto pos = ret.find('(');
  if (pos != std::string::npos) {
    auto pos2 = ret.find('+', pos);
    if (pos2 != std::string::npos) {
      auto pos3 = ret.find(')', pos2);
      auto offset = ret.substr(pos2, pos3 - pos2);
      if (pos2 != pos + 1) {
        ret = ret.substr(pos + 1, pos2 - pos - 1);
        char* demangled = abi::__cxa_demangle(ret.data(), nullptr, nullptr, nullptr);
        if (demangled != nullptr) {
          ret = demangled;
          free(demangled);
        } else {
          ret += "()";
          ret += offset;
        }
      } else {
        ret = ret.substr(0, pos) + offset;
        auto slash = ret.rfind('/');
        if (slash != std::string::npos) {
          ret = ret.substr(slash + 1);
        }
      }
    }
  }
  return ret;
}

void MemoryTracer::Reset() {
  allocated_memory_.clear();
}

void MemoryTracer::Decrese(void* ptr) {
  if (auto it = allocated_memory_.find(ptr); it != allocated_memory_.end()) {
    allocated_memory_.erase(it);
  }
}

size_t MemoryTracer::Check() {
  size_t total_leak = 0;
  for (auto& [ptr, alloc_info] : allocated_memory_) {
    total_leak += alloc_info.size_;
    LOG_PRINTLN("Memory leak: %p, size: %lu, allocated at: %s", ptr, alloc_info.size_,
                Addr2sym(alloc_info.call_).c_str());
  }

  allocated_memory_.clear();

  return total_leak;
}

#endif

void MemoryCounter::Count() {
  MemoryContextMemConsumed(CurrentMemoryContext, reinterpret_cast<MemoryContextCounters*>(this));
}

}  // namespace optimizer

namespace std::pmr {

class DefaultResourceImpl : public memory_resource {
 public:
  void* do_allocate(std::size_t bytes, std::size_t align) override {
    if (auto* ptr = palloc_aligned(bytes, static_cast<size_t>(align), MCXT_ALLOC_NO_OOM); ptr != nullptr) {
      optimizer::AddCheck(ptr, bytes);
      return ptr;
    }
    throw std::bad_alloc();
  }

  void do_deallocate(void* pointer, std::size_t /*bytes*/, std::size_t /*alignment*/) override {
    optimizer::DeleCheck(pointer);
    pfree(pointer);
  }

  [[nodiscard]] bool do_is_equal(const memory_resource& other) const noexcept override { return &other == this; }
};

memory_resource* get_default_resource() noexcept {  // NOLINT
  // Yes, this leaks. We have had SO many problems with the default memory
  // resource going out of scope before the other things were cleaned up that we
  // decided to live with the leak, rather than running into races over and over
  // again.

  // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,bugprone-unhandled-exception-at-new)
  static auto* default_resource_instance = new DefaultResourceImpl();
  return default_resource_instance;
}

memory_resource* new_delete_resource() noexcept {  // NOLINT
  return get_default_resource();
}

// NOLINTNEXTLINE: lint.sh thinks there is a C-style cast in the next line.
memory_resource* set_default_resource(memory_resource* /*resource*/) noexcept {
  return get_default_resource();
}
}  // namespace std::pmr
